from argparse import ArgumentParser
import struct
import os
from os import path


def decode(f, sz, fmt):
    s = f.read(sz)
    return struct.unpack("<" + fmt, s)


def extract(file, chip_id, chip_ver, fw_ver, extra_ver, i, j, offset, size, td):
    s = open(file, 'rb')
    s.seek(offset)
    data = s.read(size)
    s.close()
    p = path.join(td, "{}_{:x}_{:x}_{:x}p{}".format(
        i, chip_id, chip_ver, extra_ver, j))
    d = open(p, 'wb')
    d.write(data)
    d.close()


def main():
    parser = ArgumentParser()
    parser.add_argument('file')
    parser.add_argument('-e', '--extract', type=str, metavar='DIR')
    args = parser.parse_args()
    with open(args.file, 'rb') as file:
        signature = file.read(32)
        if not signature.startswith(b"MTK_DOWNLOAD_AGENT"):
            print('Not a Mediatek download agent')
            exit(1)

        da_id = file.read(0x40).rstrip(b"\0")
        print('id: %s' % str(da_id, "ascii"))
        assert decode(file, 4, "I")[0] == 4
        assert file.read(4) == b"\x99\x88\x66\x22"
        num_socs = decode(file, 4, "I")[0]
        print("num_chips:", num_socs)

        print('%5s | %20s | %20s | %20s | %20s' %
              ('#', 'Chip ID', 'Chip Ver', 'FW Ver', 'Extra Ver'))

        if args.extract:
            os.mkdir(args.extract)
        for i in range(num_socs):
            rec = file.read(0xdc)
            fields = struct.unpack("<2sHIIIIIIIQIIIIIIIIIIIII128s", rec)
            assert fields[0] == b'\xda\xda'
            fields = fields[1:-1]
            chip_id = fields[0]
            chip_ver = fields[1]
            fw_ver = fields[2]
            extra_ver = fields[3]

            print('%5d | %20x | %#20x | %#20x | %#20x' %
                  (i, chip_id, chip_ver, fw_ver, extra_ver))
            for j, no in enumerate([5, 9, 14, 19]):
                print('      %3d [%#x, 0x%05x, 0x%08x]  # %#x-%#x' % (j, fields[no],
                                                                      fields[no + 1], fields[no + 2], fields[no], fields[no] + fields[no + 1]))
                if args.extract:
                    extract(args.file, chip_id, chip_ver, fw_ver, extra_ver,
                            i, j, fields[no], fields[no+1], args.extract)

        file.close()


if __name__ == "__main__":
    main()
